import asyncio
import websockets
import json
import time
import random
import aiohttp
from aiortc import RTCPeerConnection, RTCSessionDescription
from aiortc.contrib.media import MediaPlayer, MediaRelay
from aiortc.rtcrtpsender import RTCRtpSender

# 推流参数全局变量
pc = None
relay = None
webcam = None
streaming_lock = asyncio.Lock()  # 控制推流的锁
initial_stream_started = False  # 标志是否为首次推流
video_sender = None

# 配置推流信息
ip = "39.106.77.250"
stream_url = f"webrtc://{ip}:1990/live/livestream1?schema=https"
api = f"https://{ip}:1990/rtc/v1/publish/"

async def start_stream(resolution="1920x1080", framerate="30", is_initial=False):
    global pc, relay, webcam, initial_stream_started

    async with streaming_lock:
        if pc:
            await stop_stream()

        if not is_initial:
            print("Waiting for 60 seconds before starting new stream...")
            # await asyncio.sleep(60)  # 非首次推流等待一分钟

        print(f"Starting stream with resolution {resolution} @ {framerate} fps")
        
        # 创建 RTCPeerConnection
        pc = RTCPeerConnection()

        # 设置视频选项
        options = {
            "framerate": framerate,
            "video_size": resolution,
            "input_format": "mjpeg"
        }

        # 初始化 MediaPlayer
        webcam = MediaPlayer("/dev/video0", format="v4l2", options=options)
        relay = MediaRelay()
        video_sender = pc.addTrack(relay.subscribe(webcam.video))

        # 设置视频编码为 H.264
        force_codec = "video/H264"
        codecs = RTCRtpSender.getCapabilities("video").codecs
        transceiver = next(
            t for t in pc.getTransceivers() if t.sender == video_sender
        )
        transceiver.setCodecPreferences([c for c in codecs if c.mimeType == force_codec])

        # 创建和设置本地描述
        offer = await pc.createOffer()
        await pc.setLocalDescription(offer)

        data = {
            "api": api,
            "tid": hex(int(time.time()) * random.randint(1000, 9999))[2:9],
            "streamurl": stream_url,
            "clientip": None,
            "sdp": offer.sdp
        }

        async with aiohttp.ClientSession() as session:
            async with session.post(api, data=json.dumps(data), ssl=False) as resp:
                response = await resp.json()

        await pc.setRemoteDescription(
            RTCSessionDescription(sdp=response["sdp"], type="answer")
        )
        initial_stream_started = True
        print("Stream started successfully!")

async def stop_stream():
    global pc, relay, webcam
    if pc:
        print("Stopping current stream...")
        await pc.close()  # 关闭 RTCPeerConnection
        pc = None
    if webcam:
        print("Closing MediaPlayer...")
        webcam.video.stop()
        webcam = None
    # webcam = None
    # if video_sender:
    #     print("Closing video_sender...")
    #     video_sender.stop()  # 关闭 video_sender，释放资源
    #     video_sender = None

    relay = None  # 清理 relay 对象
    print("Stream stopped.")



async def receive_commands(uri):
    async with websockets.connect(uri) as websocket:
        print("Connected to WebSocket server")
        try:
            while True:
                command = await websocket.recv()
                if command.startswith("video|"):
                    _, resolution, framerate = command.split("|")
                    print(f"Received command to change to {resolution} @ {framerate} fps")
                    if not streaming_lock.locked():
                        asyncio.create_task(start_stream(resolution, framerate))
                    else:
                        print("Stream already in progress, ignoring new command.")
        except websockets.exceptions.ConnectionClosed:
            print("WebSocket connection closed, reconnecting...")
            await asyncio.sleep(5)
            await receive_commands(uri)

async def main():
    initial_resolution = "1280x720"
    initial_framerate = "30"
    await start_stream(initial_resolution, initial_framerate, is_initial=True)
    uri = "ws://39.106.77.250:8089/ws?command_receiver"
    await receive_commands(uri)

if __name__ == "__main__":
    asyncio.run(main())
